"""
性能相关装饰器
"""

import functools
import time

from myUtils.logger import setup_logger

logger = setup_logger()


def timeit(func):
    """测量函数执行时间的装饰器"""

    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        start_time = time.perf_counter()
        result = func(*args, **kwargs)
        end_time = time.perf_counter()
        execution_time = end_time - start_time
        print('==================================================')
        logger.info(
            f"Function '{func.__name__}' took {execution_time:.4f} seconds to complete"
        )
        print('==================================================')
        return result

    return wrapper


def timeit_generator(func):
    """测量生成器函数执行时间的装饰器"""

    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        start_time = time.perf_counter()
        gen = func(*args, **kwargs)
        return timeit_gen_wrapper(gen, start_time, func.__name__)

    return wrapper


def timeit_gen_wrapper(gen, start_time, func_name):
    """生成器包装器"""
    try:
        for item in gen:
            yield item
    finally:
        end_time = time.perf_counter()
        execution_time = end_time - start_time
        print('==================================================')
        logger.info(
            f"Generator '{func_name}' took {execution_time:.4f} seconds to complete"
        )
        print('==================================================')


def timeit_async(func):
    """测量异步函数执行时间的装饰器"""

    @functools.wraps(func)
    async def wrapper(*args, **kwargs):
        start_time = time.perf_counter()
        result = await func(*args, **kwargs)
        end_time = time.perf_counter()
        execution_time = end_time - start_time
        logger.info(
            f"Async function '{func.__name__}' took {execution_time:.4f} seconds to complete"
        )
        return result

    return wrapper


def timeit_async_generator(func):
    """测量异步生成器函数执行时间的装饰器"""

    @functools.wraps(func)
    async def wrapper(*args, **kwargs):
        start_time = time.perf_counter()
        gen = func(*args, **kwargs)
        async for item in timeit_async_gen_wrapper(gen, start_time,
                                                   func.__name__):
            yield item

    return wrapper


async def timeit_async_gen_wrapper(gen, start_time, func_name):
    """异步生成器包装器"""
    try:
        async for item in gen:
            yield item
    finally:
        end_time = time.perf_counter()
        execution_time = end_time - start_time
        logger.info(
            f"Async generator '{func_name}' took {execution_time:.4f} seconds to complete"
        )
